#include "HandTrackerLKHSVCD.hpp"

#include <iostream>

using namespace cv;
using namespace std;

HandTrackerLKHSVCD::HandTrackerLKHSVCD( bool useHSV, Rect trackingRect, Ptr< GeometricGestureClassifier > gClassifier, float measurementNoiseCovScalar )
: useHSV(useHSV), trackingRect( trackingRect ), gClassifier( gClassifier ), measurementNoiseCovScalar( measurementNoiseCovScalar ) {
	
	//FIX : Test these values //FULLTEST
	maxCorners = 200;
	minCornersToTrack = 100;
	qualityLevel = 0.01;;
	minDistance = 3;
	
	shortDistIgnoreThreshold = 50.0;
	
	cornerReuseErrorThreshold = 2000;
}

AbstractHandTracker* HandTrackerLKHSVCD::init( bool	useHSV, double processNoiseCovScalar, double measurementNoiseCovScalar,
											  Rect trackingRect, Mat& initImage, Ptr< GeometricGestureClassifier > gClassifier ) {
	
	HandTrackerLKHSVCD* ht = new HandTrackerLKHSVCD( useHSV, trackingRect, gClassifier, (float) measurementNoiseCovScalar  );
	
	ht->imgPtr1 = &(ht->imgA);//TODO could maybe move this to constructor?
	ht->imgPtr2 = &(ht->imgB);
	
	if ( useHSV ) {
		cvtColor( initImage, ht->imgB, CV_BGR2HSV );
	} else {
		cvtColor( initImage, ht->imgB, CV_BGR2GRAY );
	}
	
	
	
	ht->KF.init( 2, 2 ); //Two "dynamParams": (x,y). Two "measureParams", also (x,y)	
	
	setIdentity( ht->KF.transitionMatrix );
	setIdentity( ht->KF.measurementMatrix );
	//FIXNOW REVISE VALUES
	setIdentity( ht->KF.processNoiseCov, Scalar::all(processNoiseCovScalar) ); //FIXNOW ?  //1e-1 good?
	setIdentity( ht->KF.measurementNoiseCov, Scalar::all(measurementNoiseCovScalar) ); //DIFFICULT! 1 is too sensitive and 10 is too insensitive?
	setIdentity( ht->KF.errorCovPost, Scalar::all(1) ); //TODO drop arg scalar all?
	
	//ht->KF.transitionMatrix = (Mat_<float>(2, 2) << (trackingRect.x + trackingRect.width/2), 0, 0, (trackingRect.y + trackingRect.height/2) );
	
	//FIX does this work?
	ht->KF.statePre = (Mat_<float>(2, 1) << (trackingRect.x + trackingRect.width/2), (trackingRect.y + trackingRect.height/2) );
	ht->KF.statePost = (Mat_<float>(2, 1) << (trackingRect.x + trackingRect.width/2), (trackingRect.y + trackingRect.height/2) );
	
	//NOTE: sets all elements to zero!  //= Mat::zeros( 2, 1, CV_32F);
	
	
	
	// -- visualisation -- //
	{
		namedWindow( "trackingTest", 1 );
		//namedWindow( "classFail", 1 );
		//cvMoveWindow( "classFail", 100, 700);
	}
	
	return ht;
}


Rect HandTrackerLKHSVCD::track(cv::Mat& nextImage ) {
	cout << "------------------------------" << endl;
	
	//VISUALISATION //TODO RAT : for visualisation only!
	Mat imgToDraw;
	nextImage.copyTo( imgToDraw );
	
	Rect oldTrackingRect(trackingRect);
	//VISUALISATION -- draw older tracking rects //TODO RAT
	rectangle( imgToDraw, Point( oldTrackingRect.x, oldTrackingRect.y ), Point( oldTrackingRect.x + oldTrackingRect.width, oldTrackingRect.y + oldTrackingRect.height ),
			  Scalar( 0, 0, 255) ); //NOTE: oldTrackingRect is used later for more visualisation
	
	int countOfReusedCorners;
	
	
	// -- swap pointers to prepare for next image -- //
	Mat* imgPtrTmp = imgPtr1;
	imgPtr1 = imgPtr2;
	imgPtr2 = imgPtrTmp;
	
	Mat featureImg;
	
	if ( useHSV ) {
		//  -- convert image to HSV -- //
		cvtColor( nextImage, *imgPtr2, CV_BGR2HSV );
		
		// -- set ROI for feature selection -- //
		Mat imgGrey; //TODO goodFeaturesToTrack only runs on grey? 
		cvtColor( nextImage, imgGrey, CV_BGR2GRAY );
		featureImg = Mat( imgGrey, trackingRect );
		
	} else {
		//  -- convert image to GRAY -- //
		cvtColor( nextImage, *imgPtr2, CV_BGR2GRAY );
		
		// -- set ROI for feature selection -- //
		featureImg = Mat( *imgPtr2, trackingRect );
	}
	
	// -- run classifier -- //
	
	// get hand image for classifier
	Mat handImg( nextImage, trackingRect );
	
	bool classified = gClassifier->classify( handImg );
	
	Mat contourImage = gClassifier->getContourImage();
	
	// -- Decide on corners to reuse -- //
	
	//NOTE: also generate meanshift image
	
	Mat reusedCornersMeanshiftImg = Mat::zeros( nextImage.size(), CV_8UC1 ); //TODO speed up?
	int reusedCornersMeanshiftCount = 0;
	
	// "Step 1 of reuse" 
	// These are points we previously tracked with error below the threshold cornerReuseErrorThreshold.
	// Of these, we keep only those that lie on the newly computed hand contour.
	
	if (classified) { //TODO only on this condition, yeah?
		for( int i=0; i<potentiallyReusedCorners.size(); i++ ) {
			Point corner(
						 cvRound( potentiallyReusedCorners[i].x ),
						 cvRound( potentiallyReusedCorners[i].y )
						 );
			
			//BUG: What happens when index reaches outside contour image!?
			if ( trackingRect.x <= corner.x && corner.x < trackingRect.x + trackingRect.width &&
				 trackingRect.y <= corner.y && corner.y < trackingRect.y + trackingRect.height &&
				contourImage.at<uchar>( corner.y  - trackingRect.y, corner.x  - trackingRect.x ) != 0 ) { //WARNING: since this is matrix index notation, rows go before columns, and hence y before x.
				//FIXNOW correct? CHECK BY VISUALISING LOOKUP!
				//cout << "corner " << i << " REUSED" << endl;
				reusedCorners.push_back( potentiallyReusedCorners[i] );
				reusedCornersMeanshiftImg.at<uchar>( corner ) = 255; //WARNING: when a point is input rather than two indices, x should go before y!
				reusedCornersMeanshiftCount++;
				circle(imgToDraw, corner, 1, Scalar(255, 0, 0), 2);
			}
		} 
	}//TODO update to take advantage of memory continuity?
	potentiallyReusedCorners.clear(); //done with these for now. Prepare for new corners
	
	
	if ( reusedCorners.size() > 0 ) {
		
		int msTrackStartX = trackingRect.x;
		int msTrackStartY = trackingRect.y;
		
		cout << "pre  ms X=" << trackingRect.x << " Y=" << trackingRect.y << endl;
		
		meanShift( reusedCornersMeanshiftImg, trackingRect, TermCriteria( TermCriteria::MAX_ITER + TermCriteria::EPS, 5, 0.01) ); //FULLTEST ?
		
		//FIXNOW does that make sense?
		Mat measurement = (Mat_<float>(2, 1)
						   << (KF.statePost.at<float>(0)+(trackingRect.x  - msTrackStartX)), (KF.statePost.at<float>(1) + (trackingRect.y - msTrackStartY)) );
		
		
		float meanshiftNoiseFactor = (float) minCornersToTrack / (float) reusedCornersMeanshiftCount / 4;
		
		KF.measurementNoiseCov *= meanshiftNoiseFactor; //Modify precision according to number of corners tracked!
		//FIX semantics wrong? Would want precision to depend on number of corners still in view after meanshift?
		
		
		KF.correct( measurement );
		KF.predict(); //FIXNOWs is that necessary each time?
		
		setIdentity( KF.measurementNoiseCov, Scalar::all(measurementNoiseCovScalar) );
		
		trackingRect.x = cvRound( KF.statePost.at<float>(0) ) - trackingRect.width/2;
		trackingRect.y = cvRound( KF.statePost.at<float>(1) ) - trackingRect.height/2;
		//FIXNOW Must update corner positions! No, wait, it's already in big frame?
		
		cout << "pre  ms X=" << trackingRect.x << " Y=" << trackingRect.y << endl;
		cout << "precision ms/regular=" << meanshiftNoiseFactor << endl;
		
		//forcing rectangle to stay inside frame  //TODO is this necessary?
		if ( trackingRect.x < 0 ) trackingRect.x = 0;
		if ( nextImage.cols < trackingRect.x + trackingRect.width ) trackingRect.x = nextImage.cols - trackingRect.width;
		if ( trackingRect.y < 0 ) trackingRect.y = 0;
		if ( nextImage.rows < trackingRect.y + trackingRect.height ) trackingRect.y = nextImage.rows - trackingRect.height;
		
		
		//VISUALISATION //TODO RAT
		rectangle( imgToDraw, Point( trackingRect.x, trackingRect.y ), Point( trackingRect.x + trackingRect.width, trackingRect.y + trackingRect.height ),
				  Scalar(255, 0, 0) );
	}
	
	// -- select features (corners) to track -- //
	corners.clear(); //Prepare for new corners.
	if (reusedCorners.size() < minCornersToTrack ) {
		
		
		int featureFindCount = maxCorners - reusedCorners.size();
		
		if (classified) {
			goodFeaturesToTrack( featureImg, corners, featureFindCount, qualityLevel, minDistance, gClassifier->getContourImage());
		} else {
			goodFeaturesToTrack( featureImg, corners, featureFindCount, qualityLevel, minDistance );
		}
		//FIXNOW what to do on track fail? Give up?
		//FIX dilate mask?
		
		//Change frame of reference for corners to be the full image
		for( int i=0; i<corners.size(); i++) {
			corners[i].x += trackingRect.x;
			corners[i].y += trackingRect.y;
		} //TODO update to take advantage of memory continuity?
		
		countOfReusedCorners = reusedCorners.size(); //VISUALISATION //TODO RAT : for visualisation only!
	}
	
	//	//TODO reenable
	//	//Add the reused corners
	//	for( int i=0; i<reusedCorners.size(); i++) {
	//		corners.push_back( reusedCorners[i] );
	//	} //TODO update to take advantage of memory continuity?
	
	//TODO RAT
	corners.swap(reusedCorners);
	for( int i=0; i<reusedCorners.size(); i++) {
		corners.push_back( reusedCorners[i] );
	} //Now we've put the reused corners at the front
	
	reusedCorners.clear();//done with these for now. Prepare for new corners
	
	// -- track features -- //
	
	calcOpticalFlowPyrLK( *imgPtr1, *imgPtr2, corners, nextCorners, cornersFound, cornerErrors );
	//FIX test / tweak, use derivLambda? how many levels? etc.
	
	// -- determine displacement of tracking rect -- //
	
	int trackedCornerCount;
	//TODO RAT
	bool shortDistIgnore[corners.size()];
	
	double startX = KF.statePost.at<float>(0);
	double startY = KF.statePost.at<float>(1);
	
	for( int j=0; j<corners.size(); j++ ) {
		if (cornersFound[j] == 0) { //TODO also discard if error is high?
			//cout << "corner " << j << " NOT FOUND" << endl;
			//corner not found
			continue;
		}
		
		//FIX REWRITE to make faster when done with //VISUALISATION
		float shortDistIgnoreHack = (nextCorners[j].x  - corners[j].x) * (nextCorners[j].x  - corners[j].x);
		shortDistIgnoreHack += (nextCorners[j].y  - corners[j].y) * (nextCorners[j].y  - corners[j].y);
		
		if (shortDistIgnoreHack < shortDistIgnoreThreshold) {
			//cout << "corner " << j << " TRACKED BUT IN BACKGROUND" << endl;
			shortDistIgnore[j] = true;
		} else {
			shortDistIgnore[j] = false;
			
			if ( cornerErrors[j] < cornerReuseErrorThreshold ) {
				potentiallyReusedCorners.push_back(nextCorners[j]);
				//cout << "corner " << j << " TRACKED AND CONSIDERED FOR REUSE" << endl;
			} else {
				//cout << "corner " << j << " TRACKED BUT EXCEEDS ERROR THRES" << endl;
			}
			
			
			//FIXNOW does that make sense?
			Mat measurement = (Mat_<float>(2, 1) << (startX+(nextCorners[j].x  - corners[j].x)), (startY + (nextCorners[j].y - corners[j].y)) );
			
			KF.correct( measurement );
			KF.predict(); //FIXNOWs is that necessary each time?
			
			trackedCornerCount++;
			
		} //FIXNOW: only ignore if some points do move?
	}
	
	
	if ( trackedCornerCount > 0 ) {
		
		trackingRect.x = cvRound( KF.statePost.at<float>(0) ) - trackingRect.width/2;
		trackingRect.y = cvRound( KF.statePost.at<float>(1) ) - trackingRect.height/2;
		
		//forcing rectangle to stay inside frame
		if ( trackingRect.x < 0 ) trackingRect.x = 0;
		if ( nextImage.cols < trackingRect.x + trackingRect.width ) trackingRect.x = nextImage.cols - trackingRect.width;
		if ( trackingRect.y < 0 ) trackingRect.y = 0;
		if ( nextImage.rows < trackingRect.y + trackingRect.height ) trackingRect.y = nextImage.rows - trackingRect.height;
		
	}
	
	// -- visualisation -- //  //VISUALISATION
	{
		
		//Tracking rect
		rectangle( imgToDraw, Point( trackingRect.x, trackingRect.y ), Point( trackingRect.x + trackingRect.width, trackingRect.y + trackingRect.height ),
				  Scalar(0,255,0) );
		
		if (classified) {
			//mask for LK features
			Mat maskInBGR;
			
			cvtColor( gClassifier->getContourImage(), maskInBGR, CV_GRAY2BGR );
			
			Mat imgToDraw_handRect( imgToDraw, oldTrackingRect );
			
			imgToDraw_handRect = imgToDraw_handRect + ( 0.3 * maskInBGR );
		}
		
		
		
		for( int j=0; j<corners.size(); j++ ) {
			if (cornersFound[j] == 0) { //TODO also discard if error is very high?
				//cout << "corner " << j << " REUSED NOT FOUND" << endl;
				//corner not found
				continue;
			}
			
			Scalar colour;
			
			if (shortDistIgnore[j]) {
				colour = Scalar( 0, 0, 255 ); // Scalar( 0, 0, 0 );
				
			} else if (j < countOfReusedCorners) {
				//colour = Scalar( 55 * ( 600 - cornerErrors[j] ) / 600, 200 + 55 * ( 600 - cornerErrors[j] ) / 600, 55 * ( 600 - cornerErrors[j] ) / 600);
				colour = Scalar( 255, 0, 0 ); //colour = Scalar( 200 + 55 * ( 600 - cornerErrors[j] ) / 600, 55 * ( 600 - cornerErrors[j] ) / 600, 55 * ( 600 - cornerErrors[j] ) / 600);
			} else {
				colour = Scalar( 0, 255, 0 );
			}
			
			Point p0(
					 cvRound( corners[j].x ),
					 cvRound( corners[j].y )
					 );
			Point p1(
					 cvRound( nextCorners[j].x ),
					 cvRound( nextCorners[j].y )
					 );
			line( imgToDraw, p0, p1, colour ); //draw by thickness?
		}
		
		imshow( "trackingTest", imgToDraw );
	}
	
	
	
	return trackingRect;
	
}



